import chat_model
from langchain_core.messages import HumanMessage
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

model_runable = chat_model.get_model()

#做一个chatbotmanager类，管理所有的会话
class ChatbotManager():

    store = {} # 存储所有会话的存储字典，key为session_id，value为InMemoryChatMessageHistory


    # 根据session_id获取会话历史的方法
    @classmethod
    def get_session_history(cls, session_id) -> InMemoryChatMessageHistory:
        if session_id not in cls.store:
            cls.store[session_id] = InMemoryChatMessageHistory()
        return cls.store[session_id]

#做一个chatbot类，就是一个聊天机器人
class ChatBot():

    def __init__(self, session_id: str):
        # 通过提供的session_id，实例化不同的chatbot

        # 传入session_id后，它的记忆就根据session_id得到
        self.chat_runable = RunnableWithMessageHistory(model_runable, ChatbotManager.get_session_history)
        self.config = {"configurable": {"session_id": session_id}}
        '''
        RunnableWithMessageHistory 是带消息历史的runable
        声明方式 chain = RunnableWithMessageHistory(runable, 
        根据session_id得到InMemoryChatMessageHistory()的一个方法名)

        调用方式：
        .invoke(messages列表,config = {"configurable": {"session_id": "[id变量]"}})
        若流式，则.steam
        '''

    # 单伦对话
    def chat(self, content):
        messages = [HumanMessage(content=content)]
        response = self.chat_runable.stream(
            messages, # messages列表
            config=self.config, # config
        )
        for chunk in response:
            print(chunk.content, end="")

    # 模拟多轮对话
    def play(self):
        while True:
            content = input("请输入：")
            if content == "exit":
                break
            self.chat(content)
            print()

if __name__ == "__main__":
    bot = ChatBot("test")
    bot.play()